# OCR模块 - 用于屏幕内容理解和文本提取
import os
import cv2
import numpy as np
import pyautogui
import time
import logging
import sys
from PIL import Image
import pytesseract
from typing import Dict, List, Optional, Tuple

# 配置日志，确保使用UTF-8编码
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("ocr_module.log", encoding='utf-8'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger("OCRModule")

class OCRProcessor:
    """OCR处理器，负责从屏幕截图或图像中提取文本"""
    
    def __init__(self, language: str = 'chi_sim+eng'):
        """初始化OCR处理器
        
        Args:
            language: OCR识别语言，默认为中文简体+英文
        """
        self.language = language
        self.ensure_tesseract_installed()
        logger.info(f"OCR处理器初始化完成，使用语言: {language}")
    
    def ensure_tesseract_installed(self):
        """确保Tesseract OCR已正确安装"""
        # 尝试从多个来源获取Tesseract路径
        tesseract_path = None
        
        # 1. 首先检查是否已经配置
        try:
            pytesseract.get_tesseract_version()
            current_path = pytesseract.pytesseract.tesseract_cmd
            logger.info(f"Tesseract已配置，当前路径: {current_path}")
            return
        except pytesseract.TesseractNotFoundError:
            pass
        
        # 2. 尝试从配置文件读取
        config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tesseract_config.ini')
        if os.path.exists(config_path):
            try:
                import configparser
                config = configparser.ConfigParser()
                config.read(config_path, encoding='utf-8')
                if 'Tesseract' in config and 'path' in config['Tesseract']:
                    tesseract_path = config['Tesseract']['path'].strip('"').strip("'")
                    logger.info(f"从配置文件读取Tesseract路径: {tesseract_path}")
            except Exception as e:
                logger.error(f"读取配置文件失败: {str(e)}")
        
        # 3. 尝试从环境变量读取
        if not tesseract_path:
            env_path = os.environ.get('TESSERACT_PATH')
            if env_path:
                tesseract_path = env_path.strip('"').strip("'")
                logger.info(f"从环境变量读取Tesseract路径: {tesseract_path}")
        
        # 4. 尝试Windows默认安装路径
        if not tesseract_path:
            default_paths = [
                r'C:\Program Files\Tesseract-OCR\tesseract.exe',
                r'C:\Program Files (x86)\Tesseract-OCR\tesseract.exe'
            ]
            for path in default_paths:
                if os.path.exists(path):
                    tesseract_path = path
                    logger.info(f"使用Windows默认Tesseract路径: {tesseract_path}")
                    break
        
        # 验证并设置路径
        if tesseract_path and os.path.exists(tesseract_path):
            try:
                pytesseract.pytesseract.tesseract_cmd = tesseract_path
                # 验证设置是否成功
                pytesseract.get_tesseract_version()
                logger.info("Tesseract OCR配置成功")
            except Exception as e:
                logger.error(f"Tesseract OCR配置失败: {str(e)}")
        else:
            logger.warning("Tesseract OCR未找到。请安装Tesseract OCR并确保可访问。")
            logger.warning("可以通过以下方式配置Tesseract路径:")
            logger.warning("1. 创建tesseract_config.ini文件并设置path参数")
            logger.warning("2. 设置TESSERACT_PATH环境变量")
            logger.warning("3. 安装到默认路径: C:\\Program Files\\Tesseract-OCR")
    
    def extract_text_from_image(self, image: Image.Image) -> str:
        """从图像中提取文本
        
        Args:
            image: PIL图像对象
            
        Returns:
            提取的文本
        """
        try:
            # 转换为灰度图以提高识别率
            gray_image = image.convert('L')
            
            # 应用二值化
            threshold = 128
            binary_image = gray_image.point(lambda x: 255 if x > threshold else 0, '1')
            
            # 使用Tesseract OCR提取文本
            text = pytesseract.image_to_string(binary_image, lang=self.language)
            
            logger.info(f"成功从图像中提取文本，长度: {len(text)}字符")
            return text.strip()
        except Exception as e:
            logger.error(f"OCR文本提取失败: {str(e)}")
            return ""
    
    def capture_screen_and_extract(self, region: Optional[Tuple[int, int, int, int]] = None) -> Dict:
        """截图并提取文本
        
        Args:
            region: 截图区域(x, y, width, height)，None表示全屏
            
        Returns:
            包含截图和提取文本的字典
        """
        try:
            # 截图
            screenshot = pyautogui.screenshot(region=region)
            
            # 提取文本
            text = self.extract_text_from_image(screenshot)
            
            # 保存截图
            screenshot_dir = "screenshots"
            os.makedirs(screenshot_dir, exist_ok=True)
            filename = os.path.join(screenshot_dir, f"ocr_{time.strftime('%Y%m%d_%H%M%S')}.png")
            screenshot.save(filename)
            
            logger.info(f"截图并OCR处理完成，保存到: {filename}")
            
            return {
                "success": True,
                "text": text,
                "screenshot_path": filename,
                "region": region
            }
        except Exception as e:
            logger.error(f"截图或OCR处理失败: {str(e)}")
            return {
                "success": False,
                "error": str(e),
                "text": ""
            }
    
    def find_text_position(self, target_text: str, region: Optional[Tuple[int, int, int, int]] = None) -> List[Dict]:
        """在屏幕或指定区域中查找文本位置
        
        Args:
            target_text: 要查找的文本
            region: 查找区域(x, y, width, height)，None表示全屏
            
        Returns:
            找到的文本位置列表
        """
        try:
            # 截图
            screenshot = pyautogui.screenshot(region=region)
            img_array = np.array(screenshot)
            
            # 使用Tesseract获取文本位置信息
            data = pytesseract.image_to_data(img_array, lang=self.language, output_type=pytesseract.Output.DICT)
            
            results = []
            n_boxes = len(data['text'])
            
            for i in range(n_boxes):
                text = data['text'][i].strip()
                if target_text.lower() in text.lower() and data['conf'][i] > 50:  # 置信度过滤
                    # 计算位置信息
                    left = data['left'][i]
                    top = data['top'][i]
                    width = data['width'][i]
                    height = data['height'][i]
                    confidence = data['conf'][i]
                    
                    # 如果指定了区域，调整坐标为全局坐标
                    if region:
                        left += region[0]
                        top += region[1]
                    
                    results.append({
                        "text": text,
                        "position": (left, top, left + width, top + height),
                        "confidence": confidence,
                        "center": (left + width // 2, top + height // 2)
                    })
            
            logger.info(f"在屏幕中查找文本 '{target_text}'，找到 {len(results)} 个匹配项")
            return results
        except Exception as e:
            logger.error(f"查找文本位置失败: {str(e)}")
            return []
    
    def analyze_screen_content(self, region: Optional[Tuple[int, int, int, int]] = None) -> Dict:
        """分析屏幕内容，提取结构化信息
        
        Args:
            region: 分析区域(x, y, width, height)，None表示全屏
            
        Returns:
            屏幕内容分析结果
        """
        try:
            # 获取截图和文本
            result = self.capture_screen_and_extract(region)
            
            if not result["success"]:
                return result
            
            # 分析文本结构
            lines = result["text"].split('\n')
            lines = [line.strip() for line in lines if line.strip()]
            
            # 提取可能的标题（较长或大写的行）
            possible_titles = []
            for line in lines:
                if len(line) > 10 or line.isupper():
                    possible_titles.append(line)
            
            # 提取可能的按钮文本（较短的行）
            possible_buttons = []
            for line in lines:
                if 2 <= len(line) <= 20:
                    possible_buttons.append(line)
            
            analysis = {
                "success": True,
                "text": result["text"],
                "screenshot_path": result["screenshot_path"],
                "line_count": len(lines),
                "possible_titles": possible_titles[:5],  # 最多返回5个可能的标题
                "possible_buttons": possible_buttons[:10],  # 最多返回10个可能的按钮
                "region": region
            }
            
            logger.info(f"屏幕内容分析完成，提取了 {len(lines)} 行文本")
            return analysis
        except Exception as e:
            logger.error(f"屏幕内容分析失败: {str(e)}")
            return {
                "success": False,
                "error": str(e)
            }

class ScreenInteractionManager:
    """屏幕交互管理器，基于OCR结果进行屏幕交互"""
    
    def __init__(self, ocr_processor: OCRProcessor):
        """初始化屏幕交互管理器
        
        Args:
            ocr_processor: OCR处理器实例
        """
        self.ocr_processor = ocr_processor
        logger.info("屏幕交互管理器初始化完成")
    
    def click_on_text(self, target_text: str, wait_time: float = 1.0) -> bool:
        """点击屏幕上的指定文本
        
        Args:
            target_text: 要点击的文本
            wait_time: 点击后等待时间
            
        Returns:
            是否点击成功
        """
        try:
            # 查找文本位置
            positions = self.ocr_processor.find_text_position(target_text)
            
            if not positions:
                logger.warning(f"未找到文本 '{target_text}'")
                return False
            
            # 选择置信度最高的位置
            best_match = max(positions, key=lambda x: x["confidence"])
            center_x, center_y = best_match["center"]
            
            # 移动鼠标并点击
            pyautogui.moveTo(center_x, center_y, duration=0.5)
            pyautogui.click()
            
            logger.info(f"已点击文本 '{target_text}' 在位置 ({center_x}, {center_y})")
            time.sleep(wait_time)
            return True
        except Exception as e:
            logger.error(f"点击文本 '{target_text}' 失败: {str(e)}")
            return False
    
    def read_and_respond_to_dialog(self, wait_time: float = 2.0) -> Dict:
        """读取并响应对话框
        
        Args:
            wait_time: 等待对话框出现的时间
            
        Returns:
            对话框分析和响应结果
        """
        try:
            time.sleep(wait_time)
            
            # 分析屏幕内容
            analysis = self.ocr_processor.analyze_screen_content()
            
            if not analysis["success"]:
                return analysis
            
            # 检测常见对话框模式
            dialog_type = "unknown"
            suggested_action = None
            
            text_lower = analysis["text"].lower()
            
            # 检测确认对话框
            if any(keyword in text_lower for keyword in ["确认", "确定", "continue", "ok"]):
                dialog_type = "confirmation"
                suggested_action = "click_on_text(确认)" if "确认" in text_lower else "click_on_text(OK)"
            
            # 检测错误对话框
            elif any(keyword in text_lower for keyword in ["错误", "error", "失败", "failed"]):
                dialog_type = "error"
                suggested_action = "click_on_text(确定)" if "确定" in text_lower else "click_on_text(OK)"
            
            # 检测警告对话框
            elif any(keyword in text_lower for keyword in ["警告", "warning", "注意", "attention"]):
                dialog_type = "warning"
                suggested_action = "click_on_text(确定)" if "确定" in text_lower else "click_on_text(Cancel)"
            
            # 检测输入对话框
            elif any(keyword in text_lower for keyword in ["请输入", "input", "enter"]):
                dialog_type = "input"
                suggested_action = "点击输入框并输入内容"
            
            result = {
                "success": True,
                "dialog_type": dialog_type,
                "text": analysis["text"],
                "suggested_action": suggested_action,
                "possible_buttons": analysis["possible_buttons"]
            }
            
            logger.info(f"检测到对话框类型: {dialog_type}")
            return result
        except Exception as e:
            logger.error(f"读取对话框失败: {str(e)}")
            return {
                "success": False,
                "error": str(e)
            }
    
    def monitor_screen_for_changes(self, baseline_text: str, timeout: int = 30, interval: float = 2.0) -> bool:
        """监控屏幕文本变化
        
        Args:
            baseline_text: 基线文本
            timeout: 超时时间（秒）
            interval: 检查间隔（秒）
            
        Returns:
            是否检测到变化
        """
        try:
            start_time = time.time()
            
            while time.time() - start_time < timeout:
                current_result = self.ocr_processor.capture_screen_and_extract()
                
                if current_result["success"] and current_result["text"] != baseline_text:
                    logger.info("检测到屏幕文本变化")
                    return True
                
                time.sleep(interval)
            
            logger.warning(f"在 {timeout} 秒内未检测到屏幕变化")
            return False
        except Exception as e:
            logger.error(f"监控屏幕变化失败: {str(e)}")
            return False

# 模块初始化函数
def initialize_ocr_module() -> Tuple[OCRProcessor, ScreenInteractionManager]:
    """初始化OCR模块
    
    Returns:
        OCR处理器和屏幕交互管理器实例
    """
    try:
        ocr_processor = OCRProcessor()
        screen_manager = ScreenInteractionManager(ocr_processor)
        logger.info("OCR模块初始化完成")
        return ocr_processor, screen_manager
    except Exception as e:
        logger.error(f"OCR模块初始化失败: {str(e)}")
        raise